import h5py
import matplotlib.pyplot as plt
import numpy as np


def plot_export_data:
    f = h5py.File('./datasets/export_datasets/wireframe_test.h5',"r")
    e = h5py.File('./datasets/export_datasets/wireframe_test_gt.h5',"r")

    for index in list(f.keys()):
        img = f[index]['image']
        junctions=e[index]['junctions']
        line_map = e[index]['line_map']



        plt.figure('image')
        plt.imshow(img, cmap='gray')
        #画点
        for p in junctions:
            plt.scatter(p[1],p[0],s=5,c='g')
        #划线
        for i in range(line_map.shape[0]):
            # if i==0:
            #     continue
            flag = False
            for j in range(line_map[i].shape[0]):
                if line_map[i][j]==1:
                    idx=j
                    flag=True
            if not flag:
                continue
           # idx=list(idx)[0].tolist()[0]
           # print(idx)
            plt.plot([junctions[i][1],junctions[idx][1]],[junctions[i][0],junctions[idx][0]],color='green',linewidth=1)

    # i=1
    # idx=np.where(line_map[i]==1)
    # idx=list(idx)[0].tolist()[0]
    # plt.plot([junctions[i][1],junctions[idx][1]],[junctions[1][0],junctions[idx][0]],color='green',linewidth=5)


    plt.axis('off')
    # savepath = '/home/jby/python_project/line_detection_experiment/self_attention_标签/train/' + str(index) + '.png'
    savepath = '/home/jby/python_project/line_detection_experiment/self_attention_标签/test/' + str(index) + '.png'
    plt.savefig(savepath, dpi = 1200, bbox_inches = 'tight')
    plt.clf()
    # plt.figure('raw_img')
    # plt.imshow(img, cmap='gray')
    # plt.axis('off')
    # plt.show()


# e2= h5py.File('./datasets/export_datasets/official/wireframe_train_gt.h5',"r")
# junctions2=e[index]['junctions']
# line_map2 = e[index]['line_map']
# plt.figure('image2')
# plt.imshow(img)
# #画点
# for i in junctions2:
#     plt.scatter(i[1],i[0],s=5,c='g')
# #划线
# for i in range(line_map2.shape[0]):
#     if i==0:
#         continue
#     # import ipdb
#     # ipdb.set_trace()
#     flag = False
#     for j in range(i):
#         if line_map2[i][j]==1:
#             idx=j
#             flag=True
#     if not flag:
#         continue
#    # idx=list(idx)[0].tolist()[0]
#    # print(idx)
#     plt.plot([junctions[i][1],junctions[idx][1]],[junctions[i][0],junctions[idx][0]],color='green',linewidth=1)
# plt.axis('off')
#
#
# plt.show()